#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rospy
import sys
import moveit_commander
import numpy as np
from geometry_msgs.msg import WrenchStamped ,TwistStamped
from moveit_commander.conversions import pose_to_list

# 力传感器回调函数
def force_callback(msg):
    global force_data
    force_data = np.array([msg.twist.linear.x, msg.twist.linear.y, msg.twist.linear.z])
    # rospy.loginfo(f"接收到的力传感器数据: {force_data}")

# 初始化ROS节点
rospy.init_node('admittance_control_moveit', anonymous=True)

# 初始化MoveIt
moveit_commander.roscpp_initialize(sys.argv)
robot = moveit_commander.RobotCommander()
scene = moveit_commander.PlanningSceneInterface()
group = moveit_commander.MoveGroupCommander("manipulator")

# 订阅力传感器数据
force_data = np.zeros(3)
rospy.Subscriber('/force_data', TwistStamped, force_callback)

# 顺应控制参数
M = np.diag([1.0, 1.0, 1.0])  # 质量矩阵
D = np.diag([150.0, 150.0, 150.0])  # 阻尼矩阵
K = np.diag([1000.0, 1000.0, 1000.0])  # 刚度矩阵

# 期望力
F_desired = np.zeros(3)

# 控制循环频率
rate = rospy.Rate(100)  # 100 Hz

# 顺应控制循环
while not rospy.is_shutdown():
    # 计算力误差
    F_error = F_desired - force_data

    # 顺应控制规律
    acceleration = np.linalg.inv(M) @ (F_error - D @ F_error - K @ F_error)
    velocity = np.zeros(3)
    position = np.zeros(3)

    # 更新位置、速度和加速度
    position += velocity / 100
    velocity += acceleration / 100

    # 获取当前末端执行器姿态
    current_pose = group.get_current_pose().pose
    current_pose_list = pose_to_list(current_pose)

    # 根据计算出的位置更新末端执行器姿态
    current_pose_list[0] += position[0]
    current_pose_list[1] += position[1]
    current_pose_list[2] += position[2]

    # 将更新后的姿态设置为新的目标
    group.set_pose_target(current_pose_list)

    group.go(wait=True)
    rospy.loginfo("机械臂位置已更新")

    # 以控制循环频率休眠
    rate.sleep()
